#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <mpi.h>
#include <driver_types.h>

#include "../util.h"

#define ROWSIZE 2000
#define COLSIZE 4000
#define FILENAME "../2000x4000.bin"

/* a short function to print a message and exit */
void error_exit(char msg[]) {
    fprintf(stderr, "%s", msg);
    MPI_Finalize(); exit(EXIT_FAILURE);
}

float *calculation(int task_id, char *filename, float *inAG, float *inBG, float *outG, int row, int col);

/* ---- main program ---- */

int main(int argc, char *argv[]) {

  int num_tasks;

  double start_time, end_time;

  int i;
  int task_id;
  int total_task_id = 0;
  int process_num_tasks = 0;
  int total_process_num_tasks = 0;
  int verbose = 0;
  char* usage_msg = "parameters:  num_tasks [--verbose]\n";
  int required_parms = 2;
  int nprocs;
  int process_ID;
  MPI_Status status;

  float *inAG, *inBG, *outG;
  float *out;
  int sizeIn, sizeOut;
  cudaError_t err;

  if (MPI_Init(&argc, &argv) != MPI_SUCCESS) {
      fprintf(stderr, "MPI initialization error\n"); 
      exit(EXIT_FAILURE);
  }

  if (argc < required_parms) {
      error_exit(usage_msg);
  }
  num_tasks = atoi(argv[1]);
  if (num_tasks <= 0) {
      error_exit(usage_msg);
  }
  if (argc > required_parms) {
      if (strcmp(argv[required_parms], "--verbose") == 0) 
          verbose = 1;
      else
          error_exit(usage_msg);
  }

  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &process_ID);

  if (nprocs < 2) {
      error_exit("number of processes must be at least 2\n");
  }

  MPI_Barrier(MPI_COMM_WORLD);
  start_time = MPI_Wtime();

  if (process_ID == 0) {
      for (i = 0; i < num_tasks; ++i) {

          MPI_Recv(0, 0, MPI_INT, MPI_ANY_SOURCE, 0, 
                  MPI_COMM_WORLD, &status);

          task_id = i;
          MPI_Send(&task_id, 1, MPI_INT, status.MPI_SOURCE, 0, 
                  MPI_COMM_WORLD);
      }
      task_id = -1;
      for (i = 1; i < nprocs; ++i) {
          MPI_Send(&task_id, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
      }
  }
  else { // worker
      sizeIn = ROWSIZE*COLSIZE*sizeof(float);
      sizeOut = ROWSIZE*ROWSIZE*sizeof(float);

      if ((err=cudaMalloc((void **) &inAG,sizeIn)) != cudaSuccess) {
         fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err)); 
         MPI_Abort(MPI_COMM_WORLD,1);
      }
      if ((err=cudaMalloc((void **) &inBG,sizeIn)) != cudaSuccess) {
         fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
         MPI_Abort(MPI_COMM_WORLD,1);
      }
      if ((err=cudaMalloc((void **) &outG,sizeOut)) != cudaSuccess) {
         fprintf(stderr,"cudaMalloc error: %s!\n",cudaGetErrorString(err));
         MPI_Abort(MPI_COMM_WORLD,1);
      }

      task_id = 0;
      while (task_id >= 0) {

          MPI_Send(0, 0, MPI_INT, 0, 0, MPI_COMM_WORLD);

          MPI_Recv(&task_id, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);

          if (task_id >= 0) {
              ++process_num_tasks;
              if (verbose)
                  fprintf(stderr,"(process %d) task = %d\n", process_ID, task_id);
              start_time = MPI_Wtime();
              out = calculation(task_id,FILENAME,inAG,inBG,outG,ROWSIZE,COLSIZE);
              end_time = MPI_Wtime();
              write_submatrix(out,ROWSIZE,ROWSIZE,ROWSIZE,ROWSIZE);
              free(out);
              fprintf(stderr,"(process %d) finished task = %d in %g\n", process_ID, task_id,(end_time-start_time));
          }
      }
      fprintf(stderr,"process %d number of tasks = %d\n", process_ID,
              process_num_tasks);

      cudaFree(inAG);
      cudaFree(inBG);
      cudaFree(outG);

  }

  MPI_Reduce(&process_num_tasks, &total_process_num_tasks, 1,
          MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);

  MPI_Barrier(MPI_COMM_WORLD);
  end_time = MPI_Wtime();

  if (process_ID == 0) {
      fprintf(stderr,"\nMPI parallel version with %d worker processes"
              ", dynamic assignment of tasks\n", 
              nprocs-1);
      fprintf(stderr,"number of tasks = %d\n", total_process_num_tasks);
      fprintf(stderr,"running time = %g\n", end_time - start_time);
  }

  MPI_Finalize();
  return EXIT_SUCCESS;
}

